#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/RCIn.h>
#include <tf2_ros/transform_broadcaster.h>
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>

class RCTransformNode {
private:
    ros::NodeHandle nh_;
    ros::Subscriber rc_sub_;
    ros::Subscriber pose_sub_;
    ros::Subscriber flu_cmd_sub_;
    ros::Publisher pose_pub_;
    ros::Publisher cmd_pub_;

    tf2_ros::TransformBroadcaster tf_broadcaster_;
    tf2_ros::StaticTransformBroadcaster static_tf_broadcaster_;

    geometry_msgs::PoseStamped origin_pose_;
    double origin_yaw_ = 0.0;
    bool triggered_ = false;
    bool origin_set_ = false;

public:
    RCTransformNode() {
        rc_sub_ = nh_.subscribe("/mavros/rc/in", 10, &RCTransformNode::rcCallback, this);
        pose_sub_ = nh_.subscribe("/mavros/local_position/pose", 10, &RCTransformNode::poseCallback, this);
        flu_cmd_sub_ = nh_.subscribe("/mavros/setpoint_position/local_flu", 10, &RCTransformNode::fluCmdCallback, this);

        pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("/mavros/local_position/pose_flu", 10);
        cmd_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("/mavros/setpoint_position/local", 10);

        publishStaticTransform();
    }

    void rcCallback(const mavros_msgs::RCIn::ConstPtr& msg) {
        if (!triggered_ && msg->channels.size() > 9 && msg->channels[9] > 1500) {
            ROS_INFO("RC Channel 10 > 1500, ready to record origin pose on next pose message.");
            triggered_ = true;
        }
    }

    void poseCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        if (!triggered_)
            return;

        if (!origin_set_) {
            origin_pose_ = *msg;

            tf2::Quaternion q_orig;
            tf2::fromMsg(origin_pose_.pose.orientation, q_orig);
            tf2::Matrix3x3 m(q_orig);
            double roll, pitch, yaw;
            m.getRPY(roll, pitch, yaw);
            origin_yaw_ = yaw;

            origin_set_ = true;
            ROS_INFO("Origin pose recorded with yaw: %.3f", origin_yaw_);
            return;
        }

        // 当前姿态 yaw
        tf2::Quaternion q_curr;
        tf2::fromMsg(msg->pose.orientation, q_curr);
        tf2::Matrix3x3 m_curr(q_curr);
        double roll, pitch, yaw;
        m_curr.getRPY(roll, pitch, yaw);
        double relative_yaw = yaw - origin_yaw_;

        // 坐标转换
        double dx = msg->pose.position.x - origin_pose_.pose.position.x;
        double dy = msg->pose.position.y - origin_pose_.pose.position.y;

        double x_new = cos(-origin_yaw_) * dx - sin(-origin_yaw_) * dy;
        double y_new = sin(-origin_yaw_) * dx + cos(-origin_yaw_) * dy;
        double z_new = msg->pose.position.z - origin_pose_.pose.position.z;

        tf2::Quaternion q_new;
        q_new.setRPY(roll, pitch, relative_yaw);

        geometry_msgs::PoseStamped pose_out;
        pose_out.header.stamp = msg->header.stamp;
        pose_out.header.frame_id = "map_flu";
        pose_out.pose.position.x = x_new;
        pose_out.pose.position.y = y_new;
        pose_out.pose.position.z = z_new;
        pose_out.pose.orientation = tf2::toMsg(q_new);
        pose_pub_.publish(pose_out);

        // tf map_flu -> body_flu
        geometry_msgs::TransformStamped tf;
        tf.header.stamp = ros::Time::now();
        tf.header.frame_id = "map_flu";
        tf.child_frame_id = "body_flu";
        tf.transform.translation.x = x_new;
        tf.transform.translation.y = y_new;
        tf.transform.translation.z = z_new;
        tf.transform.rotation = tf2::toMsg(q_new);
        tf_broadcaster_.sendTransform(tf);
    }

    void fluCmdCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        if (!origin_set_) return;

        double x = msg->pose.position.x;
        double y = msg->pose.position.y;
        double z = msg->pose.position.z;

        // 旋转回原坐标系
        double x_map = cos(origin_yaw_) * x - sin(origin_yaw_) * y + origin_pose_.pose.position.x;
        double y_map = sin(origin_yaw_) * x + cos(origin_yaw_) * y + origin_pose_.pose.position.y;
        double z_map = z + origin_pose_.pose.position.z;

        // yaw 处理
        tf2::Quaternion q_flu;
        tf2::fromMsg(msg->pose.orientation, q_flu);
        double roll, pitch, yaw_flu;
        tf2::Matrix3x3(q_flu).getRPY(roll, pitch, yaw_flu);
        double yaw_map = yaw_flu + origin_yaw_;

        tf2::Quaternion q_map;
        q_map.setRPY(roll, pitch, yaw_map);

        geometry_msgs::PoseStamped cmd_out;
        cmd_out.header.stamp = msg->header.stamp;
        cmd_out.header.frame_id = "map";
        cmd_out.pose.position.x = x_map;
        cmd_out.pose.position.y = y_map;
        cmd_out.pose.position.z = z_map;
        cmd_out.pose.orientation = tf2::toMsg(q_map);
        cmd_pub_.publish(cmd_out);
    }

    void publishStaticTransform() {
        geometry_msgs::TransformStamped static_tf;
        static_tf.header.stamp = ros::Time::now();
        static_tf.header.frame_id = "body_flu";
        static_tf.child_frame_id = "camera_flu";
        static_tf.transform.translation.x = 0.0;
        static_tf.transform.translation.y = 0.0;
        static_tf.transform.translation.z = 0.0;
        static_tf.transform.rotation.x = 0.0;
        static_tf.transform.rotation.y = 0.0;
        static_tf.transform.rotation.z = 0.0;
        static_tf.transform.rotation.w = 1.0;

        static_tf_broadcaster_.sendTransform(static_tf);
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "rc_pose_transform_node");
    RCTransformNode node;
    ros::spin();
    return 0;
}